{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Redacting Personally Identifiable Information with NVIDIA's GLiNER-PII Model\n",
    "\n",
    "This notebook demonstrates the use of NVIDIA's [GLiNER PII model](https://huggingface.co/nvidia/gliner-PII) in NeMo Curator. The GLiNER PII model detects and classifies a broad range of Personally Identifiable Information (PII) and Protected Health Information (PHI) in structured and unstructured text. It is non-generative and produces span-level entity annotations with confidence scores across 55+ categories.\n",
    "\n",
    "This tutorial requires at least 1 NVIDIA GPU with:\n",
    "  - Volta™ or higher (compute capability 7.0+)\n",
    "  - CUDA 12.x\n",
    "\n",
    "Before running this notebook, see this [Installation Guide](https://docs.nvidia.com/nemo/curator/latest/admin/installation.html#admin-installation) page for instructions on how to install NeMo Curator. Be sure to use an installation method which includes GPU dependencies.\n",
    "\n",
    "Additionally, the tutorial uses the [GLiNER](https://github.com/urchade/GLiNER) library. It can be installed with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !uv pip install gliner"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Silence Curator logs via Loguru\n",
    "import os\n",
    "\n",
    "os.environ[\"LOGURU_LEVEL\"] = \"ERROR\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following imports are required for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from gliner_pii_redactor import GlinerPiiRedactor\n",
    "\n",
    "from nemo_curator.core.client import RayClient\n",
    "from nemo_curator.pipeline import Pipeline\n",
    "from nemo_curator.stages.text.io.reader.jsonl import JsonlReader\n",
    "from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a pipeline in NeMo Curator, we must start a Ray cluster. This can be done manually (see the [Ray documentation](https://docs.ray.io/en/latest/ray-core/starting-ray.html)) or with Curator's `RayClient`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client = RayClient()\n",
    "    ray_client.start()\n",
    "except Exception as e:\n",
    "    msg = f\"Error initializing Ray client: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Read, PII Redaction, and Write Stages\n",
    "\n",
    "Functions in NeMo Curator are called stages. For this tutorial, we will initialize 3 stages: a JSONL file reader, a PII redactor stage, and a JSONL file writer.\n",
    "\n",
    "For this tutorial, we will create a sample JSONL file to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_file_path = \"./input_data_dir\"\n",
    "\n",
    "# Create sample dataset for the tutorial\n",
    "text = [\n",
    "    \"Please contact John Doe at johndoe@example.com for more information.\",\n",
    "    \"Jane's address is 123 Main St, Anytown, USA. Jane likes it there.\",\n",
    "    \"The phone number of the company is (123) 456-7890.\",\n",
    "    \"The patient is 45 years old and has a ssn of 123-45-6789.\",\n",
    "    \"This text does not contain any sensitive information.\",\n",
    "]\n",
    "df = pd.DataFrame({\"text\": text})\n",
    "\n",
    "try:\n",
    "    os.makedirs(input_file_path, exist_ok=True)\n",
    "    df.to_json(input_file_path + \"/data.jsonl\", orient=\"records\", lines=True)\n",
    "except Exception as e:\n",
    "    msg = f\"Error creating input file: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define the reader stage with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read existing directory of JSONL files\n",
    "read_stage = JsonlReader(input_file_path, files_per_partition=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the `GlinerPiiRedactor` as a wrapper around the [GLiNER](https://github.com/urchade/GLiNER) library. Please see the `gliner_pii_redactor.py` script for the full `ProcessingStage` implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the Gliner PII Redactor\n",
    "pii_redactor_stage = GlinerPiiRedactor(text_field=\"text\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can define a stage for writing the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write results to a directory\n",
    "output_file_path = \"./gliner_pii_redactor_results\"\n",
    "\n",
    "# Use mode=\"overwrite\" to overwrite the output directory if it already exists\n",
    "# This helps to ensure that the correct output is written\n",
    "write_stage = JsonlWriter(output_file_path, mode=\"overwrite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Pipeline\n",
    "\n",
    "In NeMo Curator, we use pipelines to run distributed data workflows using Ray. Pipelines take care of resource allocation and autoscaling to achieve enhanced performance and minimize GPU idleness.\n",
    "\n",
    "For PII redaction, we are able to achieve speedups by ensuring that model inference is run in parallel across all available GPUs, while other stages such as I/O are run across all available CPUs. This is possible because Curator pipelines are composable, which allows each stage in a pipeline to run independently and with its own specified hardware resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pipeline(name='pii_redaction_pipeline', stages=[jsonl_reader(JsonlReader), gliner_pii_redactor(GlinerPiiRedactor), jsonl_writer(JsonlWriter)])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pii_redaction_pipeline = Pipeline(name=\"pii_redaction_pipeline\", description=\"Run a PII redaction pipeline\")\n",
    "\n",
    "# Add stages to the pipeline\n",
    "pii_redaction_pipeline.add_stage(read_stage)\n",
    "pii_redaction_pipeline.add_stage(pii_redactor_stage)\n",
    "pii_redaction_pipeline.add_stage(write_stage)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Composability is also what allows a model to sit between pre-processing and post-processing stages. Typical text pre-processing add-ons include text normalization (lowercasing, URL/email removal, Unicode cleanup) and language identification and filtering (to keep only target languages). A full pipeline may look something like:\n",
    "\n",
    "```python\n",
    "pipeline = Pipeline(name=\"full_pipeline\")\n",
    "pipeline.add_stage(read_stage)                # reader (JSONL/S3/etc.)\n",
    "pipeline.add_stage(lang_id_stage)             # optional: language filter\n",
    "pipeline.add_stage(pii_redactor_stage)        # pii model and redactor\n",
    "pipeline.add_stage(write_stage)               # writer (JSONL/Parquet)\n",
    "```\n",
    "\n",
    "# Run the PII Redaction Pipeline\n",
    "\n",
    "Let's run the full pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-10-28 11:57:04,806\tINFO worker.py:1692 -- Using address 127.0.1.1:6379 set in the environment variable RAY_ADDRESS\n",
      "2025-10-28 11:57:04,809\tINFO worker.py:1833 -- Connecting to existing Ray cluster at address: 127.0.1.1:6379...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2025-10-28 11:57:04,831\tINFO usage_lib.py:447 -- Usage stats collection is disabled.\n",
      "2025-10-28 11:57:04,831\tINFO scripts.py:914 -- \u001b[37mLocal node IP\u001b[39m: \u001b[1m127.0.1.1\u001b[22m\n",
      "2025-10-28 11:57:06,802\tSUCC scripts.py:950 -- \u001b[32m--------------------\u001b[39m\n",
      "2025-10-28 11:57:06,803\tSUCC scripts.py:951 -- \u001b[32mRay runtime started.\u001b[39m\n",
      "2025-10-28 11:57:06,803\tSUCC scripts.py:952 -- \u001b[32m--------------------\u001b[39m\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:954 -- \u001b[36mNext steps\u001b[39m\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:957 -- To add another node to this Ray cluster, run\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:960 -- \u001b[1m  ray start --address='127.0.1.1:6379'\u001b[22m\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:969 -- To connect to this Ray cluster:\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:971 -- \u001b[35mimport\u001b[39m\u001b[26m ray\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:972 -- ray\u001b[35m.\u001b[39m\u001b[26minit(_node_ip_address\u001b[35m=\u001b[39m\u001b[26m\u001b[33m'127.0.1.1'\u001b[39m\u001b[26m)\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:984 -- To submit a Ray job using the Ray Jobs CLI:\n",
      "2025-10-28 11:57:06,803\tINFO scripts.py:985 -- \u001b[1m  RAY_API_SERVER_ADDRESS='http://127.0.0.1:8265' ray job submit --working-dir . -- python my_script.py\u001b[22m\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:994 -- See https://docs.ray.io/en/latest/cluster/running-applications/job-submission/index.html \n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:998 -- for more information on submitting Ray jobs to the Ray cluster.\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1003 -- To terminate the Ray runtime, run\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1004 -- \u001b[1m  ray stop\u001b[22m\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1007 -- To view the status of the cluster, use\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1008 --   \u001b[1mray status\u001b[22m\u001b[26m\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1012 -- To monitor and debug Ray, view the dashboard at \n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1013 --   \u001b[1m127.0.0.1:8265\u001b[22m\u001b[26m\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1020 -- \u001b[4mIf connection to the dashboard fails, check your firewall settings and network configuration.\u001b[24m\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1121 -- \u001b[36m\u001b[1m--block\u001b[22m\u001b[39m\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1122 -- This command will now block forever until terminated by a signal.\n",
      "2025-10-28 11:57:06,804\tINFO scripts.py:1125 -- Running subprocesses are monitored and a message will be printed if any of them terminate unexpectedly. Subprocesses exit with SIGTERM will be treated as graceful, thus NOT reported.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-10-28 11:57:07,196\tINFO worker.py:2004 -- Connected to Ray cluster. View the dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265 \u001b[39m\u001b[22m\n",
      "/raid/syurick/conda/ray-curator/lib/python3.12/site-packages/ray/_private/worker.py:2052: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0\n",
      "  warnings.warn(\n",
      "2025-10-28 11:57:08,157\tINFO worker.py:1692 -- Using address 127.0.1.1:6379 set in the environment variable RAY_ADDRESS\n",
      "2025-10-28 11:57:08,159\tINFO worker.py:1833 -- Connecting to existing Ray cluster at address: 127.0.1.1:6379...\n",
      "2025-10-28 11:57:08,160\tINFO worker.py:1851 -- Calling ray.init() again after it has already been called.\n",
      "Fetching 13 files: 100%|██████████| 13/13 [00:00<00:00, 4051.56it/s]\n",
      "\u001b[36m(Stage 02 - GlinerPiiRedactor pid=4084898)\u001b[0m Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.\n"
     ]
    }
   ],
   "source": [
    "# Run the pipeline\n",
    "result = pii_redaction_pipeline.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the pipeline ran to completion and the result was written to a JSONL file, we can shut down the Ray cluster with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client.stop()\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(f\"Error stopping Ray client: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect the Output\n",
    "\n",
    "The write stage returns a list of written files. We can read the output file as a Pandas DataFrame for inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Please contact {first_name} {last_name} at {em...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>{first_name}'s address is {street_address}, {c...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>The phone number of the company is {phone_numb...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>The patient is {age} years old and has a ssn o...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>This text does not contain any sensitive infor...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                text\n",
       "0  Please contact {first_name} {last_name} at {em...\n",
       "1  {first_name}'s address is {street_address}, {c...\n",
       "2  The phone number of the company is {phone_numb...\n",
       "3  The patient is {age} years old and has a ssn o...\n",
       "4  This text does not contain any sensitive infor..."
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# For simplicity, we take the first written file from the writer stage\n",
    "# In real pipelines, the writer may return multiple files (shards) or objects\n",
    "result_file = result[0].data[0]\n",
    "\n",
    "result_df = pd.read_json(result_file, lines=True)\n",
    "result_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the sensitive information was redacted as expected."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
